'''
Created on Sep 17, 2010

@author: oabalbin
'''
'''
Created on Sep 3, 2010

@author: oabalbin
'''
import os
import sys
import subprocess
import pysam
import numpy as np

from optparse import OptionParser
from collections import deque, defaultdict
from bx.intervals.intersection import Interval, IntervalTree

from Bio.Seq import Seq
from Bio import SeqIO
from Bio.SeqRecord import SeqRecord
from Bio.Alphabet import generic_dna

'''
list of changes:
    09-21-10: It does not do interval tree and list intersections. In order to do that ==> need to re-modified the location in create BAM file. 
    Because in this case uses the int.value to print and in this modification was change to intv.chr as in miyInterval class 
'''


class myInterval:
    def __init__(self,chr, start, end, name, score, strand):
        self.chr, self.start, self.end, self.name, self.score, self.strand = \
        chr, int(start), int(end), name, score, strand
        


def get_flow_cell_info(sampleInfoFile,flowCellInfo):
    
    sampleType,tissueType,QC=flowCellInfo[0],flowCellInfo[1],flowCellInfo[2]
    sampleDict = defaultdict(list)    
    for line in sampleInfoFile:
        fields = line.strip('\n').split('\t')
        if sampleType==fields[4] and tissueType==fields[5] and QC==fields[12]:
            sampleDict[fields[0]] = fields[1:]
    
    return sampleDict


def extract_readsfrom_BAM_file(mysam_file, mysam_outfile, bed):

    # go to flow cell dir
    # This can maybe rewrite using the pysam library.
    # This will allow to make quality control on the fly.
    ### Also Check waht you need to do in order to output the arf and the fastaq file for mirDeep from this point
    ### So you do not need to use the bwa_converter after going again over the files. 
    
    quality_threshold=20
    #mysam_file = pysam.Samfile(flow_cell_bam_file+'tophat_aligned_reads.bam','rb')
    reads_quality_scores, reads_trimming_length=[],[]
    
    # histogram information
    mybins=np.arange(0,41,1)  
      
    # Check the header for TopHat @PG    ID:TopHat    VN:1.0.13
    if mysam_file.header['PG']['ID']=='TopHat' and mysam_file.header['PG']['VN']=='1.0.13':
    
        for alignedread in mysam_file.fetch(bed):        
            myletter_annotations={}
            # Note that the sam file says to store the base quality score as ASCII-33.
            myletter_annotations['phred_quality'] = [ord(letter)-33 for letter in list(alignedread.qual)]            
            mean_read_quality = np.mean(myletter_annotations['phred_quality'])
            reads_quality_scores.append(mean_read_quality)
            
            if mean_read_quality >= quality_threshold:
                # Print the read in the mirDeep2 format.
                # Print in a fastq format as required for mirDeep
                mysam_outfile.write(alignedread)
                
                
        quality_histogram = np.histogram(np.array(reads_quality_scores),bins=mybins)                
    
    else:
        # Need to see what to do with other versions.
        print mysam_file.header['PG']
        
            
    # go back to analysis directory
    #os.chdir(analysis_dir_path)
    
    return quality_histogram


def create_interval_tree(bed_file):
    
    intervals_bed_file = open(bed_file)
    intersecter = IntervalTree()
    
    for line in intervals_bed_file:
        fields = line.strip('\n').split('\t')
        nfields = len(fields)
        if nfields < 3:
            print "The bed file has not enough information"
        elif nfields == 6:
            chr_tmp=fields[0].replace('chr','')
            name, score, strand=fields[3],fields[4],default_strand(fields[5])
            intersecter.insert_interval( Interval( int(fields[1]), int(fields[2]), value={'chr':chr_tmp, 'name':name,'score':score}, strand=strand ) )
        elif nfields < 6:
            chr_tmp=fields[0].replace('chr','')
            name, score, strand="noname","noscore","+"
            intersecter.insert_interval( Interval( int(fields[1]), int(fields[2]), value={'chr':chr_tmp,'name':name, 'score':score}, strand=strand ) )
    
    intervals_bed_file.close()                
    return intersecter

def default_strand(x):
    if not x:
        print x
        return '+'
    else:
        return x


def create_list_of_intervals(intervals_bed_file):
    
    intervals_bed_file = open(intervals_bed_file)
    interval_list=deque()
    for line in intervals_bed_file:
        fields = line.strip('\n').split('\t')
        nfields = len(fields)
        if nfields < 3:
            print "The bed file has not enough information"
        elif nfields == 6:
            chr_tmp=fields[0].replace('chr','')
            name, score, strand=fields[3],fields[4],default_strand(fields[5])
            interval_list.append((myInterval(chr_tmp, int(fields[1]), int(fields[2]), name, score, strand)))
        elif nfields < 6:
            chr_tmp=fields[0].replace('chr','')
            name, score, strand="noname","noscore","+"
            interval_list.append((myInterval(chr_tmp, int(fields[1]), int(fields[2]), name, score, strand)))
    
    intervals_bed_file.close()
    return interval_list


# 1083    chr1    65296704        65296779        hsa-mir-101-1   480     -       0       0       miRNA
# create a list of mirnas intervals
def create_list_of_mirnas_intervals(intervals_bed_file):
    
    intervals_bed_file = open(intervals_bed_file)
    interval_list=deque()
    for line in intervals_bed_file:
        fields = line.strip('\n').split('\t')
        nfields = len(fields)
        if nfields < 3:
            print "The bed file has not enough information"
        elif nfields == 10 and fields[9]=='miRNA':
            chr_tmp=fields[1].replace('chr','')
            name, score, strand=fields[4]+'|'+'MIR',fields[5],default_strand(fields[6])
            interval_list.append((myInterval(chr_tmp, int(fields[2]), int(fields[3]), name, score, strand)))
    
    intervals_bed_file.close()
    return interval_list

    
def intersect_intervals_tree(list_of_intervals, interval_tree):
    
    intersected_intervals=[]
    for thisintv in list_of_intervals:
        intesected_intervals= interval_tree.find(thisintv.start,thisintv.end)
        # Specific for the repeat program. It could be eliminated from other scripts, AluSg7|Alu|SINE, MIR|MIR|SINE
        repeat_type = thisintv.name.split('|')[1]
        if repeat_type == "Alu": #"MIR": #"Alu": #MIR
            for intv in intesected_intervals:
                if (intv.value['chr']==thisintv.chr) and (intv.strand == thisintv.strand) and (intv.start < thisintv.start and intv.end > thisintv.end):
                    #print intv
                    intersected_intervals.append(intv)
                    #print thisintv.chr, thisintv.start, thisintv.end, thisintv.strand
                    #print intv, thisintv.name
    
    return intersected_intervals
    
        
def create_SAM_summary(data_flow_cell_path, analysis_dir_path, path_to_samtools, myintervals, flow_cell_dictionary, intervals_file, quality_files):
    
    outfile_intv=open(intervals_file, 'w')
    quality_file=open(quality_files, 'w')
    # array to save the quality histogram of all reads. 
    quality_histogram = np.zeros(40)
    flow_quality_histogram = np.zeros(40)
    #
    intv2write=set()
    
    for flow_cell in flow_cell_dictionary:
        
        flow_cell_path=data_flow_cell_path+flow_cell+'/'
        #print flow_cell_path
        
        flow_cell_name = flow_cell_path.split('/')[-2]
        outfile_name = analysis_dir_path+flow_cell_name+'_tu_compilation.sam'
        print outfile_name
        # Open a SAM file for reading and one for writting
        mysam_infile = pysam.Samfile(flow_cell_path+'tophat_aligned_reads.bam','rb')
        mysam_outfile = pysam.Samfile(outfile_name,'wh', header=mysam_infile.header)
        
        # Write the header of the SAM file.  
        for intv in myintervals:
            
            if len(intv.chr.split('_')) >1:
                print intv.name, intv.chr, intv.start, intv.end
                continue
            location = 'chr'+intv.chr+':'+str(intv.start)+'-'+str(intv.end)
            intv2write.add('chr'+intv.chr+'\t'+str(intv.start)+'\t'+str(intv.end))
            
            intv_quality_histogram = extract_readsfrom_BAM_file(mysam_infile, mysam_outfile, location)
            flow_quality_histogram += intv_quality_histogram[0]

        #print flow_quality_histogram
        
        quality_histogram += flow_quality_histogram
        
        mysam_infile.close()
        mysam_outfile.close()
    
    quality_file.write(",".join(map(str,quality_histogram)).replace(',','\t')+'\n')    
    
    for location in intv2write:
        outfile_intv.write(location+'\n')
        #print location
        
    outfile_intv.close()    
    
if __name__ == '__main__':
    
    
    optionparser = OptionParser("usage: %prog [options] ")
    
    optionparser.add_option("-t", "--transcripts", dest="transcripts",
                            help="Bed file with the location of transcripts of interest")    
    optionparser.add_option("-d", "--outfolder", dest="outfolder",
                            help="output folder")
    
    (options, args) = optionparser.parse_args()
    
    
    sampleInfoFile = open('/data/projects/mirnas/sampleinfo_2010_Aug_08.txt')
    data_cell_path='/exds/tuxedo/version_003_2010_Mar_26/data/'
    #data_cell_path='/exds/solexa/illumina/idea/BC_cellLinesData/PE_RNAseq/BAMs/'
    path_to_samtools='/exds/sw/bioinfo/alignment/samtools/current/samtools'
    
    transcripts=options.transcripts
    analysis_dir_path=options.outfolder
    flowCellInfo = ['Tissue','Prostate','PASS']
    intervals_file, quality_files=analysis_dir_path+'intevals_file.bed',analysis_dir_path+'quality_file.qc'
   
    #==== main =====#
    flow_cell_dictionary = get_flow_cell_info(sampleInfoFile,flowCellInfo)
    #trans_tree1 = create_interval_tree(transcripts)
    interval_list = create_list_of_intervals(transcripts)
    #interval_list = create_list_of_mirnas_intervals(mirna_interval_file)
    #intersected_intervals = intersect_intervals_tree(interval_list, trans_tree1)
    create_SAM_summary(data_cell_path, analysis_dir_path, path_to_samtools, interval_list, flow_cell_dictionary, intervals_file, quality_files)
